
import flair
from flair.data import Sentence
from flair.models import SequenceTagger
from ner_model_npu import *
import torch
import torch_npu
import time

# 检查 NPU 是否可用
if torch.npu.is_available():
    print("NPU is available.")
    flair.device = torch.device('npu')
else:
    print("NPU is not available.")
    flair.device = torch.device('cpu')

# 加载命名实体识别模型和情感分析模型
start_time=time.time()
tagger = SequenceTaggerForNPU.load_onnx("./models/en-ner-conll03-v0.4.onnx")
print("模型加载用时：",time.time()-start_time)
# 检查模型使用的设备
device = next(tagger.parameters()).device
print(f"Model is currently using device: {device}")

# 输入语句
print('Please enter the sentence to be analyzed:')
inputs = input()

# NER
tagger_sentence = Sentence(inputs)
start_time=time.time()
tagger.predict_onnx(tagger_sentence)
print("模型推理用时：",time.time()-start_time)
print(tagger_sentence.labels)



